#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <stdlib.h>
#include <errno.h>
#include <pthread.h>
#include <netinet/in.h>

#include "commtraff_socket_server.h"
#include "commtraff.h"
#include "plugin.h"

// 定义常量替换魔法数字
#define SOCKET_SEND_RETRY_COUNT 3
#define SOCKET_RETRY_DELAY_US 1000

#define BUFFER_SIZE 1024 * 1024


static pthread_t g_serverThread = 0;
static char *g_buffer = NULL;
static SocketServerCb g_cb;
static int g_serverFd = -1;
static volatile int g_serverRunning = 0;


void *ServerThreadFunc(void *arg)
{
    struct sockaddr_in serverAddr, clientAddr;
    socklen_t clientAddrLen;

    // Create UDP socket
    g_serverFd = socket(AF_INET, SOCK_DGRAM, 0);
    if (g_serverFd == -1) {
        ERROR("Failed to create UDP socket: %s", strerror(errno));
        return NULL;
    }

    // 设置端口复用
    int opt = 1;
    if (setsockopt(g_serverFd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        ERROR("Failed to set SO_REUSEADDR option: %s", strerror(errno));
        close(g_serverFd);
        g_serverFd = -1;
        return NULL;
    }

    // Set server address
    memset(&serverAddr, 0, sizeof(serverAddr));
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_addr.s_addr = INADDR_ANY;
    serverAddr.sin_port = htons(COLLECTD_SERVER_PORT);

    int recvSize = 0;
    int sendSize = 0;
    socklen_t optLen = sizeof(sendSize);
    if (getsockopt(g_serverFd, SOL_SOCKET, SO_RCVBUF, &recvSize, &optLen) < 0) {
        ERROR("get socket SO_RCVBUF errno=%d, msg=%s", errno, strerror(errno));
    }
    if (getsockopt(g_serverFd, SOL_SOCKET, SO_SNDBUF, &sendSize, &optLen) < 0) {
        ERROR("get socket SO_SNDBUF errno=%d, msg=%s", errno, strerror(errno));
    }
    INFO("Server recvSize: %d sendSize: %d", recvSize, sendSize);

    // Bind socket
    if (bind(g_serverFd, (struct sockaddr *)&serverAddr, sizeof(serverAddr)) == -1) {
        ERROR("Failed to bind socket, errno=%d, msg=%s", errno, strerror(errno));
        close(g_serverFd);
        g_serverFd = -1;
        return NULL;
    }
    g_serverRunning = 1;
    while (g_serverRunning) {
        clientAddrLen = sizeof(clientAddr);
        memset(g_buffer, 0, BUFFER_SIZE);

        // Receive data with timeout
        ssize_t bytesRead = recvfrom(g_serverFd, g_buffer, BUFFER_SIZE - 1, 0,
                                      (struct sockaddr *)&clientAddr, &clientAddrLen);

        if (bytesRead > 0) {
            if (g_cb.OnRecv) {
                g_cb.OnRecv(g_buffer, bytesRead);
            }
        } else if (bytesRead < 0) {
            if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK) {
                // 被中断或超时，继续循环检查退出条件
                continue;
            }
            ERROR("Failed to receive data: %zd, errno=%d, msg=%s", bytesRead, errno, strerror(errno));
            break;
        }
    }
    if (g_serverFd != -1) {
        close(g_serverFd);
        g_serverFd = -1;
    }
    INFO("ServerThreadFunc exit");
    return NULL;
}


int32_t SocketServerInit(SocketServerCb *cb)
{
    g_buffer = (char *)malloc(BUFFER_SIZE);
    if (g_buffer == NULL) {
        ERROR("Failed to allocate memory for buffer");
        return COMMTRAFF_FAIL;
    }
    memset(g_buffer, 0, BUFFER_SIZE);
    memset(&g_cb, 0, sizeof(g_cb));
    g_cb = *cb;
    return pthread_create(&g_serverThread, NULL, ServerThreadFunc, NULL);
}

// 添加服务器去初始化函数
int32_t SocketServerDeinit()
{
    // 设置退出标志
    g_serverRunning = 0;
    
    // 等待线程结束
    if (g_serverThread != 0) {
        pthread_join(g_serverThread, NULL);
        g_serverThread = 0;
    }
    
    if (g_buffer) {
        free(g_buffer);
        g_buffer = NULL;
    }
    
    return COMMTRAFF_SUCCESS;
}

/**
 * 向客户端发送数据的函数
 * @param clientPort 客户端端口
 * @param buffer 要发送的数据
 * @param length 数据长度
 * @return 成功返回COMMTRAFF_SUCCESS，失败返回COMMTRAFF_FAIL
 */
int32_t SendDataToClient(int32_t clientPort, uint8_t *buffer, int32_t length)
{
    if (buffer == NULL) {
        ERROR("buffer is null");
        return COMMTRAFF_FAIL;
    }

    struct sockaddr_in clientAddr;

    // Check if server socket is initialized
    if (g_serverFd == -1) {
        ERROR("Server socket not initialized");
        return COMMTRAFF_FAIL;
    }

    // Set client address
    memset(&clientAddr, 0, sizeof(clientAddr));
    clientAddr.sin_family = AF_INET;
    clientAddr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);  // localhost
    clientAddr.sin_port = htons(clientPort);

    // Send data to client using sendto with improved retry mechanism
    ssize_t bytesSent = -1;
    int retryTimes = SOCKET_SEND_RETRY_COUNT;
    do {
        bytesSent = sendto(g_serverFd, buffer, length, 0,
                            (struct sockaddr *) &clientAddr, sizeof(clientAddr));
        if (bytesSent < 0) {
            // 扩展重试错误码范围，添加重试间隔
            if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR || 
                errno == ENOBUFS || errno == ENOMEM) {
                if (retryTimes > 0) {
                    usleep(SOCKET_RETRY_DELAY_US);  // 重试间隔
                }
            } else {
                break;  // 不可恢复的错误，直接退出重试循环
            }
        }
        retryTimes--;
    } while (bytesSent < 0 && retryTimes > 0);

    if (bytesSent < 0) {
        if (errno == EAGAIN || errno == EWOULDBLOCK) {
            ERROR("Socket would block, try again later");
        } else {
            ERROR("Failed to send data to client: %s", strerror(errno));
        }
        return COMMTRAFF_FAIL;
    }

    return COMMTRAFF_SUCCESS;
}